#include "md5.h"
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdio.h>

#define F(X,Y,Z) (((X)&(Y))|((~(X))&(Z)))
#define G(X,Y,Z) (((X)&(Z))|((Y)&(~(Z))))
#define H(X,Y,Z) ((X)^(Y)^(Z))
#define I(X,Y,Z) ((Y)^((X)|(~(Z))))
#define CSHL32(x,y) (((x)<<(y))|((x)>>(32-(y))))

static uint32_t T[64];

void md5_init(void)
{
    for (size_t i = 0; i < 64; ++i)
    {
        T[i] = (uint32_t)((double)(((uint64_t)1) << 32) * fabs(sin(i + 1.0)));
    }
}

block_t *md5_make_blocks(uint64_t *num_blocks, const uint8_t *message, uint64_t message_bytes)
{
    uint64_t message_bits = message_bytes * 8;
    *num_blocks = (message_bits + BLOCK_BITS - 448) / BLOCK_BITS + 1;
    block_t *blocks = (block_t *)malloc(*num_blocks * sizeof(block_t));
    for (uint64_t i = 0; i < *num_blocks; ++i)
    {
        memset(blocks + i, 0, sizeof(blocks[0]));
    }

    uint64_t cnt = 0;
    int stop = 0;
    for (uint64_t i = 0; i < *num_blocks; ++i)
    {
        for (size_t j = 0; j < BLOCK_WORDS; ++j)
        {
            for (size_t k = 0; k < 4; ++k)
            {
                uint32_t value = message[cnt];
                if (cnt++ == message_bytes)
                {
                    value = (1 << 7);
                    stop = 1;
                }
                blocks[i].words[j] |= (value << (k * 8));
                if (stop)
                {
                    break;
                }
            }
            if (stop)
            {
                break;
            }
        }
        if (stop)
        {
            break;
        }
    }

    uint64_t ones = (((uint64_t)1) << 32) - 1;
    uint32_t lo = (uint32_t)(message_bits & ones);
    uint32_t hi = (uint32_t)(message_bits >> 32);
    blocks[*num_blocks - 1].words[BLOCK_WORDS - 2] = lo;
    blocks[*num_blocks - 1].words[BLOCK_WORDS - 1] = hi;
    return blocks;
}

void md5(uint8_t *digest, const uint8_t *message, uint64_t message_bytes)
{
    /* build blocks */
    uint64_t num_blocks;
    block_t *blocks = md5_make_blocks(&num_blocks, message, message_bytes);

    /* initialize MD buffer */
    static uint32_t P[4];
    P[0] = 0x67452301;
    P[1] = 0xefcdab89;
    P[2] = 0x98badcfe;
    P[3] = 0x10325476;

    /* main process */
    static uint32_t Q[4];
    size_t k, start;
    for (uint64_t i = 0; i < num_blocks; ++i)
    {
        const uint32_t *X = blocks[i].words;
        memcpy(Q, P, sizeof(Q));
        start = 0;
#define GQ(x) (Q[(start+x)&3])

        /* round 1 */
        k = 0;
        static const uint32_t shl1[4] = {7, 12, 17, 22};
        for (int j = 0; j < 16; ++j)
        {
            Q[start] = GQ(1) + CSHL32((GQ(0) + F(GQ(1), GQ(2), GQ(3)) + X[k] + T[j]), shl1[j & 3]);
            start = ((start + 3) & 3);
            ++k;
        }

        /* round 2 */
        k = 1;
        static const uint32_t shl2[4] = {5, 9, 14, 20};
        for (int j = 16; j < 32; ++j)
        {
            Q[start] = GQ(1) + CSHL32((GQ(0) + G(GQ(1), GQ(2), GQ(3)) + X[k] + T[j]), shl2[j & 3]);
            start = ((start + 3) & 3);
            k = (k + 5) & 15;
        }

        /* round 3 */
        k = 5;
        static const uint32_t shl3[4] = {4, 11, 16, 23};
        for (int j = 32; j < 48; ++j)
        {
            Q[start] = GQ(1) + CSHL32((GQ(0) + H(GQ(1), GQ(2), GQ(3)) + X[k] + T[j]), shl3[j & 3]);
            start = ((start + 3) & 3);
            k = (k + 3) & 15;
        }

        /* round 4 */
        k = 0;
        static const uint32_t shl4[4] = {6, 10, 15, 21};
        for (int j = 48; j < 64; ++j)
        {
            Q[start] = GQ(1) + CSHL32((GQ(0) + I(GQ(1), GQ(2), GQ(3)) + X[k] + T[j]), shl4[j & 3]);
            start = ((start + 3) & 3);
            k = (k + 7) & 15;
        }

        /* addition */
        for (int j = 0; j < 4; ++j)
        {
            P[j] += Q[j];
        }
    }

    for (int i = 0; i < 4; ++i)
    {
        for (int j = 0; j < 4; ++j)
        {
            digest[(i << 2) | j] = ((P[i] >> (j * 8)) & 255);
        }
    }
    free(blocks);
}

#define IPAD 0x36
#define OPAD 0x5c
#define MAX(x,y) ((x)>(y)?(x):(y))
void md5_hmac_md5(uint8_t *hmac, const uint8_t *message, uint64_t message_bytes, const uint8_t *key,
                  uint64_t key_bytes)
{
    if (key_bytes > BLOCK_BYTES)
    {
        memset(hmac, 0, 16);
        return;
    }
    uint8_t *tem = (uint8_t *)malloc(BLOCK_BYTES + MAX(message_bytes, 16));
    memcpy(tem, key, key_bytes);
    memset(tem + key_bytes, 0, BLOCK_BYTES - key_bytes);
    for (size_t i = 0; i < BLOCK_BYTES; ++i)
    {
        tem[i] ^= IPAD;
    }
    memcpy(tem + BLOCK_BYTES, message, message_bytes);
    md5(tem + BLOCK_BYTES, tem, BLOCK_BYTES + message_bytes);

    memcpy(tem, key, key_bytes);
    memset(tem + key_bytes, 0, BLOCK_BYTES - key_bytes);
    for (size_t i = 0; i < BLOCK_BYTES; ++i)
    {
        tem[i] ^= OPAD;
    }
    md5(hmac, tem, BLOCK_BYTES + 16);

    free(tem);
}

